#!/usr/bin/ruby

require 'optparse'
require 'ostruct'

def usage() 
   puts "Usage: udpstat [-6] [-KB|-MB] <interval> <count>"
   puts
   puts "   Options:"
   puts "     -6  - IPv6 statistics (by default IPv4 is used)"
   puts "     -KB - show utilization in KB"
   puts "     -MB - show utilization in MB"
   puts 
   puts "   <interval> - interval in which statistics are to be reported"
   puts "   <count>    - the number of statistics to be shown"
   puts 
   puts "@2012 Krystian Brachmanski"
end

sysconfig=false

interval=-1
repeat=-1
divider=1
port="[A-Z0-9]+"

options = OpenStruct.new
options.descr = "B"
options.encoding = "utf8"
options.transfer_type = :auto
options.ipv6 = false
options.pid=-1

opts = OptionParser.new do |opts|
   opts.banner = "Usage: udpstat [options] <interval> <count>"

   opts.separator ""
   opts.separator "Specific options:"

   # Optional argument; multi-line description.
   opts.on("-6", "--ipv6",
              "Enable IPv6 stat parsing") do
      options.ipv6 = true
   end

   opts.on("-KB", "--kilobytes ",
              "Present values in kilobytes") do
      options.descr = "KB"
   end

   opts.on("-MB", "--megabytes ",
              "Present values in megabytes") do
      options.descr = "MB"
   end

   opts.on("-p","--ppid pid",
	      "PID of the process to be monitored") do |pid|
      options.pid = pid
   end

   opts.separator ""
   opts.separator "Common options:"

   opts.on("-h", "--help", "Show this message") do
      puts opts
      exit
   end

   opts.on("-v","--version", "Show version") do
      puts OptionParser::Version.join('.')
      exit
   end

   opts.separator ""
   opts.separator "   <interval> - interval in which statistics are to be reported"
   opts.separator "   <count>    - the number of statistics to be shown"
   opts.separator ""

   opts.separator "@2013 Krystian Brachmanski, Wroclaw" 

end
opts.parse!(ARGV)

if ARGV.length < 2
  puts opts
  exit
end

ARGV.each { |argument|
  if interval<0
     interval=argument.to_i
  else
     repeat=argument.to_i
  end	
}

if interval==-1 or repeat==-1
  usage
  exit 1
end

path_prefix="/proc/"
if options.pid != "-1"
  puts "Monitoring pid : " + options.pid 
  path_prefix += options.pid + "/"
end

if !options.ipv6
   GLOBALUDPSTAT=path_prefix+"net/udp"
else
   GLOBALUDPSTAT=path_prefix+"net/udp6"
end

case options.descr
  when "B"
     divider=1
  when "KB"
     divider=1024
  when "MB"
     divider=1048576
end

rmem_buffer=`sysctl -n net.core.rmem_max 2>/dev/null`.to_i / divider
tmem_buffer=`sysctl -n net.core.wmem_max 2>/dev/null`.to_i / divider

#                    1946: 00000000:BD67 00000000:0000 07 00000000:00000000 00:00000000 00000000     0        0 12967 2 0000000000000000 0
line_re=Regexp.new(" *[0-9]+\: +([0-9A-Z]+)\:("+port+") +([0-9A-Z]+)\:([0-9A-Z]+) [0-9]+ +([0-9]+)\:([0-9]+) +[0-9]+\:[0-9]+ +[0-9]+ +[0-9]+ +[0-9]+ +[0-9]+ +[0-9]+ +[0-9]+ +([0-9])+")
#LINE = Regexp.new('pat', Regexp::IGNORECASE | Regexp::MULTILINE)

def convertIP(iphex)
   return iphex.byteslice(1..2).to_i(16).to_s + "." + iphex.byteslice(3..4).to_i(16).to_s + "." + iphex.byteslice(5..6).to_i(16).to_s + "." + iphex.byteslice(7..8).to_i(16).to_s
end


printf "%12s %15s%-5s %15s:%-5s	%9s %9s %9s %9s %7s", "Timestamp", "Local IP:", "port", "Remote IP", "port", "rx" + "[" + options.descr + "]", "rbuf" + "[" + options.descr + "]", "tx" + "[" + options.descr  + "]", "tbuf" + "[" + options.descr + "]", "drops"
puts

1.upto(repeat) { |i|
   sleep interval 
   t = Time.now.to_i.to_s
   File.open(GLOBALUDPSTAT).each { |line|
      m=line.match(line_re)
      if m
	if !options.ipv6
	  localIP=convertIP(m[1])
	  remoteIP= convertIP(m[3])
	else
	  localIP=m[1]
	  remoteIP=m[3]
	end
        printf "%12s %15s:%-5s %15s:%-5s %8s %8s %8s %8s %7s\n", t, localIP, m[2].to_i(16), remoteIP, m[4].to_i(16), m[6].to_i(16)/divider,rmem_buffer, m[5].to_i(16)/divider,tmem_buffer, m[7]
	puts
      end
   }
} 
